# -*- coding:utf-8 -*-
import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix
import itertools
import numpy as np

def plot_confusion_matrix(cm, classes,
                          title='Confusion matrix2',
                          cmap=plt.cm, fig_name = 'Confusion matrix.png'):
    cm = cm.astype('float') / cm.sum(axis=1)[:, np.newaxis]
    # cmap  'YlGnBu'   'viridis'  'Wistia' 'OrRd'
    cmap = plt.cm.get_cmap('YlGnBu')
    plt.imshow(cm, interpolation='nearest',cmap=cmap)
    plt.colorbar()
    # plt.title(title)
    # plt.colorbar()
    tick_marks = np.arange(len(classes))
    # tick_marks = [1, 2, 3, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 18, 21, 22, 23, 24, 25, 26]
    # classes = [1, 2, 3, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 18, 21, 22, 23, 24, 25, 26]
    plt.xticks(tick_marks, classes, rotation=45)
    plt.yticks(tick_marks, classes,rotation=45)

    thresh = cm.max() / 2.
    for i, j in itertools.product(range(cm.shape[0]), range(cm.shape[1])):
        if cm[i, j] <= 0.00001:
            cm[i, j] = 0
            plt.text(j, i, '0'.format(cm[i, j]), horizontalalignment="center",fontsize=10,
                     color="white" if cm[i, j] > thresh else "black")
        else:
            plt.text(j, i, '{:.2f}'.format(cm[i, j]), horizontalalignment="center",fontsize=10,
                     color="white" if cm[i, j] > thresh else "black")
    plt.tight_layout()
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    # plt.savefig('confusion_matrix.png', dpi=300)  # 保存图片
    #
    plt.savefig(fig_name, dpi=300, bbox_inches='tight')
    print("混淆矩阵保存成功")
    plt.show()


def read_write_txt(file = 'pre_txt.txt'):
    file_cont = open(file)
    # print(file_cont)
    for line in file_cont:
        txt_pred= line
    file_cont.close()
    # print(txt_pred)
    txt_pred = txt_pred.split(',')[0:-1]
    # print()
    # print(txt_pred)
    txt_pred = [int(i) for i in txt_pred]
    # print(txt_pred)
    return txt_pred

# 标签，0:Surprise, 1:Fear, 2:Disgust, 3:Happiness, 4:Sadness, 5:Anger, 6:Neutral
if __name__ == '__main__':
    # read_txt_comp(true_labels)
    txt_pred = read_write_txt(file = 'rafdb_pred.txt')
    true_labels = read_write_txt(file = 'rafdb_truelabel.txt')
    conf_mat = confusion_matrix(y_true=true_labels, y_pred=txt_pred)
    plot_confusion_matrix(conf_mat, classes = ['Surprise', 'Fear', 'Disgust', 'Happiness', 'Sadness', 'Anger', 'Neutral'],
                          title='',
                          cmap=plt.cm, fig_name='Confusion rafdbmatrix_txt.png')